package com.boot.utils;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.function.Function;
import java.util.stream.Collectors;

public class BatchUtil {
    /**
     * 分批次处理的每批次的大小
     */
    public static int batchSize = 200;

    /**
     * 异步执行任务的线程池
     */
    public static ExecutorService executorService =
            new ThreadPoolExecutor(0, 10, 10, TimeUnit.SECONDS, new LinkedBlockingQueue<>(2048), new ThreadFactory() {
                @Override
                public Thread newThread(Runnable r) {
                    int i = (int) (Math.random() * 100);
                    return new Thread(r, "BatchUtil-executorService-" + i);
                }
            }, new ThreadPoolExecutor.CallerRunsPolicy());

    /**
     * 同步分批次处理
     *
     * @param collect
     * @param function
     * @param <T>
     * @param <R>
     * @return
     */
    public static <T, R> List<R> batchApply(List<T> collect, Function<List<T>, List<R>> function) {
        return batchApply(collect, function, batchSize);
    }

    /**
     * 异步分批次处理
     *
     * @param collect
     * @param function
     * @param <T>
     * @param <R>
     * @return
     */
    public static <T, R> List<R> batchApplyAsync(List<T> collect, Function<List<T>, List<R>> function) {
        return batchApplyAsync(collect, function, batchSize);
    }

    /**
     * 同步分批次处理
     *
     * @param collect
     * @param function
     * @param batchSize
     * @param <T>
     * @param <R>
     * @return
     */
    public static <T, R> List<R> batchApply(List<T> collect, Function<List<T>, List<R>> function, int batchSize) {
        if (batchSize <= 0) {
            batchSize = BatchUtil.batchSize;
        }
        if (collect.size() < batchSize) {
            return function.apply(collect);
        } else {
            List<List<R>> lists = applySync(collect, function, batchSize);
            return lists.stream()
                    .flatMap(e -> e.stream())
                    .collect(Collectors.toList());
        }
    }

    /**
     * 异步分批次处理
     *
     * @param collect
     * @param function
     * @param batchSize
     * @param <T>
     * @param <R>
     * @return
     */
    public static <T, R> List<R> batchApplyAsync(List<T> collect, Function<List<T>, List<R>> function, int batchSize) {
        if (batchSize <= 0) {
            batchSize = BatchUtil.batchSize;
        }
        if (collect.size() < batchSize) {
            return function.apply(collect);
        } else {

            List<Future<List<R>>> futures = applyAsync(collect, function, batchSize);

            List<R> results = new ArrayList<>();
            for (Future<List<R>> future : futures) {
                try {
                    results.addAll(future.get());
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            }
            return results;
        }
    }


    private static <T, R> List<Future<List<R>>> applyAsync(List<T> collect, Function<List<T>, List<R>> function, int batchSize) {
        int size = collect.size();
        int start = 0;
        List<Future<List<R>>> futures = new ArrayList<>();
        while (start < size) {
            int end = start + batchSize;
            if (end > size) {
                end = size;
            }
            List<T> ts = collect.subList(start, end);
            Future<List<R>> submit = executorService.submit(() -> {
                List<R> apply = function.apply(ts);
                return apply;
            });
            futures.add(submit);
            start = start + batchSize;
        }
        return futures;
    }


    private static <T, R> List<List<R>> applySync(List<T> collect, Function<List<T>, List<R>> function, int batchSize) {
        int size = collect.size();
        int start = 0;
        List<List<R>> futures = new ArrayList<>();
        while (start < size) {
            int end = start + batchSize;
            if (end > size) {
                end = size;
            }
            List<T> ts = collect.subList(start, end);
            List<R> apply = function.apply(ts);
            futures.add(apply);
            start = start + batchSize;
        }
        return futures;
    }

}
